#include "viml/register_allocator.h"
#include "viml/instruction.h"
#include "viml/value.h"
#include <algorithm>
#include <set>

namespace viml {

bool RegisterAllocator::allocate(std::shared_ptr<Bytecode>& bytecode, const LivenessAnalyzer& analyzer) {
    if (!bytecode) {
        return false;
    }
    
    reset();
    
    // 阶段1: 使用贪心着色算法分配寄存器
    allocate_registers(analyzer);
    
    // 阶段2: 根据映射表重写字节码
    rewrite_bytecode(bytecode);
    
    // 计算重用次数
    size_t virtual_register_count = register_mapping_.size();
    if (virtual_register_count > physical_register_count_) {
        reuse_count_ = virtual_register_count - physical_register_count_;
    }
    
    return true;
}

void RegisterAllocator::reset() {
    register_mapping_.clear();
    physical_register_count_ = 0;
    reuse_count_ = 0;
}

void RegisterAllocator::allocate_registers(const LivenessAnalyzer& analyzer) {
    // TODO: 实现贪心着色算法
    //
    // 算法:
    // 1. 获取所有生命周期区间,按长度降序排列
    // 2. 对每个虚拟寄存器r:
    //    a. 获取与r冲突的所有寄存器
    //    b. 找到这些冲突寄存器已分配的物理寄存器集合
    //    c. 选择最小未使用的物理寄存器编号分配给r
    // 3. 记录映射关系
    
    const auto& intervals = analyzer.get_live_intervals();
    const auto& conflict_graph = analyzer.get_conflict_graph();
    
    // 按生命周期长度排序(优先分配生命周期长的寄存器)
    std::vector<LiveInterval> sorted_intervals = intervals;
    std::sort(sorted_intervals.begin(), sorted_intervals.end(), 
              [](const LiveInterval& a, const LiveInterval& b) {
                  return (a.end_pc - a.start_pc) > (b.end_pc - b.start_pc);
              });
    
    // 跟踪已使用的物理寄存器
    std::set<uint32_t> used_physical_regs;
    
    // 为每个虚拟寄存器分配物理寄存器
    for (const auto& interval : sorted_intervals) {
        uint32_t virtual_reg = interval.register_id;
        
        // 如果已经分配过,跳过
        if (register_mapping_.count(virtual_reg)) {
            continue;
        }
        
        // 获取冲突的寄存器集合
        std::unordered_set<uint32_t> conflicts;
        if (conflict_graph.count(virtual_reg)) {
            conflicts = conflict_graph.at(virtual_reg);
        }
        
        // 找到可用的物理寄存器
        uint32_t physical_reg = find_available_register(virtual_reg, conflicts, analyzer);
        
        // 记录映射
        register_mapping_[virtual_reg] = physical_reg;
        used_physical_regs.insert(physical_reg);
    }
    
    // 更新物理寄存器计数
    physical_register_count_ = used_physical_regs.size();
}

uint32_t RegisterAllocator::find_available_register(
    uint32_t virtual_reg,
    const std::unordered_set<uint32_t>& conflicts,
    const LivenessAnalyzer& analyzer) {
    
    // 收集冲突寄存器已分配的物理寄存器
    std::unordered_set<uint32_t> used_physical_regs;
    for (uint32_t conflict_reg : conflicts) {
        if (register_mapping_.count(conflict_reg)) {
            used_physical_regs.insert(register_mapping_[conflict_reg]);
        }
    }
    
    // 找到最小可用的物理寄存器编号
    uint32_t physical_reg = 0;
    while (used_physical_regs.count(physical_reg)) {
        physical_reg++;
    }
    
    return physical_reg;
}

void RegisterAllocator::rewrite_bytecode(std::shared_ptr<Bytecode>& bytecode) {
    // 遍历字节码,替换寄存器操作数
    size_t inst_count = bytecode->instruction_count();
    
    for (size_t pc = 0; pc < inst_count; ++pc) {
        Instruction instr = bytecode->get_instruction(pc);
        OpCode op = instr.opcode();
        bool modified = false;
        
        // 处理不同指令的寄存器操作数
        uint8_t new_a = instr.a();
        uint8_t new_b = instr.b();
        uint8_t new_c = instr.c();
        
        // 映射A寄存器
        if (register_mapping_.count(instr.a())) {
            new_a = register_mapping_[instr.a()];
            modified = true;
        }
        
        // 映射B寄存器(对于ABC格式指令)
        switch (op) {
            case OpCode::MOVE:
            case OpCode::LOAD_LOCAL:
            case OpCode::STORE_LOCAL:
            case OpCode::NEG:
            case OpCode::NOT:
            case OpCode::LEN:
            case OpCode::ITER_INIT:
            case OpCode::LIST_APPEND:
            case OpCode::ADD:
            case OpCode::SUB:
            case OpCode::MUL:
            case OpCode::DIV:
            case OpCode::MOD:
            case OpCode::EQ:
            case OpCode::NE:
            case OpCode::LT:
            case OpCode::LE:
            case OpCode::GT:
            case OpCode::GE:
            case OpCode::AND:
            case OpCode::OR:
            case OpCode::GET_INDEX:
            case OpCode::SET_INDEX:
            case OpCode::DICT_SET:
            case OpCode::STR_CAT:
            case OpCode::CALL:
                if (register_mapping_.count(instr.b())) {
                    new_b = register_mapping_[instr.b()];
                    modified = true;
                }
                break;
            default:
                break;
        }
        
        // 映射C寄存器(对于三操作数指令)
        switch (op) {
            case OpCode::ADD:
            case OpCode::SUB:
            case OpCode::MUL:
            case OpCode::DIV:
            case OpCode::MOD:
            case OpCode::EQ:
            case OpCode::NE:
            case OpCode::LT:
            case OpCode::LE:
            case OpCode::GT:
            case OpCode::GE:
            case OpCode::AND:
            case OpCode::OR:
            case OpCode::GET_INDEX:
            case OpCode::SET_INDEX:
            case OpCode::DICT_SET:
            case OpCode::STR_CAT:
                if (register_mapping_.count(instr.c())) {
                    new_c = register_mapping_[instr.c()];
                    modified = true;
                }
                break;
            default:
                break;
        }
        
        // 如果修改了,更新指令
        if (modified) {
            Instruction new_instr(op, new_a, new_b, new_c);
            bytecode->set_instruction(pc, new_instr);
        }
    }
}

} // namespace viml
